import os
import sys
import time
import asyncio
import subprocess
import logging
import pyautogui as pg
import pyperclip as pc
import uiautomation as ua
from typing import Literal, List, Optional
from contextlib import asynccontextmanager
from mcp.server.fastmcp import FastMCP

# 设置Windows控制台UTF-8编码
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')

# 配置日志
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# 设置pyautogui选项
pg.FAILSAFE = False
pg.PAUSE = 0.5

# 简单的Desktop类实现
class SimpleDesktop:
    def __init__(self):
        self.encoding = 'utf-8'
    
    def get_windows_version(self):
        try:
            result = subprocess.run(['ver'], shell=True, capture_output=True, text=True)
            return result.stdout.strip()
        except Exception:
            return "Windows"
    
    def get_default_language(self):
        try:
            import winreg
            key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r'Control Panel\\International')
            value, _ = winreg.QueryValueEx(key, 'LocaleName')
            winreg.CloseKey(key)
            return value
        except Exception:
            return "zh-CN"
    
    def launch_app(self, name: str):
        try:
            # 使用start命令启动应用
            subprocess.Popen(['start', '', name], shell=True)
            return f"已启动应用: {name}", 0
        except Exception as e:
            return f"启动应用失败: {str(e)}", 1
    
    def execute_command(self, command: str):
        try:
            # 使用PowerShell执行命令
            result = subprocess.run(['powershell', '-Command', command], 
                                  capture_output=True, text=True, timeout=60)
            return result.stdout, result.returncode
        except subprocess.TimeoutExpired:
            return "命令执行超时", 1
        except Exception as e:
            return f"命令执行失败: {str(e)}", 1
    
    def get_element_under_cursor(self):
        try:
            return ua.ControlFromCursor()
        except Exception:
            # 返回一个简单的控制对象
            class DummyControl:
                Name = "未知元素"
                ControlTypeName = "未知类型"
            return DummyControl()
    
    def resize_app(self, size=None, loc=None):
        try:
            # 获取当前活动窗口
            window = ua.GetForegroundControl()
            if size:
                window.Resize(size[0], size[1])
            if loc:
                window.Move(loc[0], loc[1])
            return "窗口大小/位置已调整", 0
        except Exception as e:
            return f"调整窗口失败: {str(e)}", 1
    
    def switch_app(self, name: str):
        try:
            # 查找包含指定名称的窗口
            windows = ua.GetRootControl().GetChildren()
            for window in windows:
                if name.lower() in window.Name.lower():
                    window.SetFocus()
                    return f"已切换到应用: {name}", 0
            return f"未找到应用窗口: {name}", 1
        except Exception as e:
            return f"切换应用失败: {str(e)}", 1
    
    def get_state(self, use_vision=False):
        # 简单实现，返回基本状态信息
        class DummyState:
            class TreeState:
                def interactive_elements_to_string(self):
                    return "简单实现: 未获取交互式元素"
                def informative_elements_to_string(self):
                    return "简单实现: 未获取信息元素"
                def scrollable_elements_to_string(self):
                    return "简单实现: 未获取可滚动元素"
            
            tree_state = TreeState()
            screenshot = None
            
            def apps_to_string(self):
                return "简单实现: 未获取已打开应用列表"
            
            def active_app_to_string(self):
                try:
                    return ua.GetForegroundControl().Name
                except Exception:
                    return "未知"
        
        return DummyState()

# 初始化Desktop对象
desktop = SimpleDesktop()

# 初始化FastMCP服务器
mcp = FastMCP(name='desktop-mcp')

# 启动/关闭钩子
@asynccontextmanager
async def lifespan(app: FastMCP):
    """启动前初始化和关闭后清理"""
    try:
        logger.info("桌面控制服务初始化中...")
        await asyncio.sleep(1)
        yield
    finally:
        logger.info("桌面控制服务正在关闭...")

# 设置生命周期钩子
mcp.lifespan = lifespan

# 启动应用工具
@mcp.tool()
def launch_tool(name: str) -> str:
    """
    从Windows开始菜单启动应用程序
    参数：
    name: 应用程序名称（如 "notepad", "calculator", "chrome"）
    """
    try:
        response, status = desktop.launch_app(name.lower())
        logger.info(response)
        return response
    except Exception as e:
        logger.error(f"启动应用失败: {str(e)}")
        return f"错误: {str(e)}"

# PowerShell命令工具
@mcp.tool()
def powershell_tool(command: str) -> str:
    """
    执行PowerShell命令并返回输出和状态码
    参数：
    command: 要执行的PowerShell命令
    """
    try:
        response, status_code = desktop.execute_command(command)
        logger.info(f"PowerShell命令执行结果: 状态码={status_code}")
        return f'Response: {response}\nStatus Code: {status_code}'
    except Exception as e:
        logger.error(f"执行PowerShell命令失败: {str(e)}")
        return f"错误: {str(e)}"

# 桌面状态工具
@mcp.tool()
def state_tool(use_vision: bool = False) -> dict:
    """
    捕获桌面状态信息
    参数：
    use_vision: 是否包含视觉截图
    """
    try:
        desktop_state = desktop.get_state(use_vision=use_vision)
        interactive_elements = desktop_state.tree_state.interactive_elements_to_string()
        informative_elements = desktop_state.tree_state.informative_elements_to_string()
        scrollable_elements = desktop_state.tree_state.scrollable_elements_to_string()
        apps = desktop_state.apps_to_string()
        active_app = desktop_state.active_app_to_string()
        
        result = {
            "default_language": desktop.get_default_language(),
            "focused_app": active_app,
            "opened_apps": apps,
            "interactive_elements": interactive_elements,
            "informative_elements": informative_elements,
            "scrollable_elements": scrollable_elements
        }
        
        logger.info(f"获取桌面状态: 当前活动应用={active_app}")
        return result
    except Exception as e:
        logger.error(f"获取桌面状态失败: {str(e)}")
        return {"error": str(e)}

# 剪贴板工具
@mcp.tool()
def clipboard_tool(mode: Literal['copy', 'paste'], text: str = None) -> str:
    """
    复制文本到剪贴板或获取当前剪贴板内容
    参数：
    mode: 'copy' 模式用于复制，'paste' 模式用于获取
    text: 在'copy'模式下要复制的文本
    """
    try:
        if mode == 'copy':
            if text:
                pc.copy(text)  # 复制文本到系统剪贴板
                logger.info(f"已复制文本到剪贴板: {text[:30]}...")
                return f'已复制 "{text}" 到剪贴板'
            else:
                raise ValueError("未提供要复制的文本")
        elif mode == 'paste':
            clipboard_content = pc.paste()  # 从系统剪贴板获取文本
            logger.info(f"已获取剪贴板内容: {clipboard_content[:30]}...")
            return f'剪贴板内容: "{clipboard_content}"'
        else:
            raise ValueError('无效模式。请使用 "copy" 或 "paste"。')
    except Exception as e:
        logger.error(f"剪贴板操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 点击工具
@mcp.tool()
def click_tool(loc: List[int], button: Literal['left', 'right', 'middle'] = 'left', clicks: int = 1) -> str:
    """
    在特定坐标点击UI元素
    参数：
    loc: 坐标 [x, y]
    button: 鼠标按钮 'left', 'right', 'middle'
    clicks: 点击次数 1, 2, 3
    """
    try:
        if len(loc) != 2:
            raise ValueError("位置必须是包含两个整数的列表 [x, y]")
        x, y = loc[0], loc[1]
        pg.moveTo(x, y)
        pg.click(x=x, y=y, button=button, clicks=clicks)
        
        num_clicks = {1: '单击', 2: '双击', 3: '三击'}
        button_names = {'left': '左键', 'right': '右键', 'middle': '中键'}
        
        logger.info(f"在坐标({x},{y})执行{num_clicks.get(clicks)}操作")
        return f'{num_clicks.get(clicks)} {button_names.get(button)} 在坐标 ({x},{y})。'
    except Exception as e:
        logger.error(f"点击操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 输入文本工具
@mcp.tool()
def type_tool(loc: List[int], text: str, clear: bool = False, press_enter: bool = False) -> str:
    """
    在输入字段中输入文本
    参数：
    loc: 目标元素坐标 [x, y]
    text: 要输入的文本
    clear: 是否清除现有文本
    press_enter: 是否按下回车键
    """
    try:
        if len(loc) != 2:
            raise ValueError("位置必须是包含两个整数的列表 [x, y]")
        x, y = loc[0], loc[1]
        pg.click(x=x, y=y)
        
        if clear:
            pg.hotkey('ctrl', 'a')
            pg.press('backspace')
        
        pg.typewrite(text, interval=0.1)
        
        if press_enter:
            pg.press('enter')
        
        logger.info(f"在坐标({x},{y})输入文本: {text[:30]}...")
        return f'已在坐标 ({x},{y}) 输入文本: {text}'
    except Exception as e:
        logger.error(f"输入文本失败: {str(e)}")
        return f"错误: {str(e)}"

# 调整窗口大小工具
@mcp.tool()
def resize_tool(size: Optional[List[int]] = None, loc: Optional[List[int]] = None) -> str:
    """
    调整活动应用程序窗口大小或移动位置
    参数：
    size: 窗口大小 [宽度, 高度]
    loc: 窗口位置 [x, y]
    """
    try:
        if size is not None and len(size) != 2:
            raise ValueError("大小必须是包含两个整数的列表 [宽度, 高度]")
        if loc is not None and len(loc) != 2:
            raise ValueError("位置必须是包含两个整数的列表 [x, y]")
        
        response, _ = desktop.resize_app(tuple(size) if size else None, 
                                        tuple(loc) if loc else None)
        
        logger.info(f"调整窗口大小/位置: 大小={size}, 位置={loc}")
        return response
    except Exception as e:
        logger.error(f"调整窗口失败: {str(e)}")
        return f"错误: {str(e)}"

# 切换应用工具
@mcp.tool()
def switch_tool(name: str) -> str:
    """
    切换到特定应用程序窗口
    参数：
    name: 应用程序名称
    """
    try:
        response, status = desktop.switch_app(name)
        logger.info(response)
        return response
    except Exception as e:
        logger.error(f"切换应用失败: {str(e)}")
        return f"错误: {str(e)}"

# 滚动工具
@mcp.tool()
def scroll_tool(loc: Optional[List[int]] = None, 
              scroll_type: Literal['horizontal', 'vertical'] = 'vertical',
              direction: Literal['up', 'down', 'left', 'right'] = 'down',
              wheel_times: int = 1) -> str:
    """
    在特定坐标或当前鼠标位置滚动
    参数：
    loc: 滚动坐标 [x, y]，默认为当前鼠标位置
    scroll_type: 'horizontal' 或 'vertical'
    direction: 滚动方向
    wheel_times: 滚动次数
    """
    try:
        if loc:
            if len(loc) != 2:
                raise ValueError("位置必须是包含两个整数的列表 [x, y]")
            x, y = loc[0], loc[1]
            pg.moveTo(x, y)
        
        if scroll_type == 'vertical':
            if direction == 'up':
                pg.scroll(wheel_times)
            elif direction == 'down':
                pg.scroll(-wheel_times)
            else:
                return '无效方向。请使用 "up" 或 "down"。'
        elif scroll_type == 'horizontal':
            # 水平滚动使用Shift+滚轮
            pg.keyDown('shift')
            time.sleep(0.05)
            if direction == 'left':
                pg.scroll(wheel_times)
            elif direction == 'right':
                pg.scroll(-wheel_times)
            else:
                return '无效方向。请使用 "left" 或 "right".'
            time.sleep(0.05)
            pg.keyUp('shift')
        else:
            return '无效类型。请使用 "horizontal" 或 "vertical".'
        
        logger.info(f"执行{scroll_type}方向{direction}滚动{wheel_times}次")
        return f'已{scroll_type}方向{direction}滚动{wheel_times}次。'
    except Exception as e:
        logger.error(f"滚动操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 拖拽工具
@mcp.tool()
def drag_tool(from_loc: List[int], to_loc: List[int]) -> str:
    """
    从源坐标拖拽到目标坐标
    参数：
    from_loc: 起始坐标 [x, y]
    to_loc: 目标坐标 [x, y]
    """
    try:
        if len(from_loc) != 2:
            raise ValueError("from_loc必须是包含两个整数的列表 [x, y]")
        if len(to_loc) != 2:
            raise ValueError("to_loc必须是包含两个整数的列表 [x, y]")
        
        x1, y1 = from_loc[0], from_loc[1]
        x2, y2 = to_loc[0], to_loc[1]
        
        # 执行拖拽操作
        pg.moveTo(x1, y1)
        pg.mouseDown()
        time.sleep(0.1)
        pg.moveTo(x2, y2, duration=0.5)
        time.sleep(0.1)
        pg.mouseUp()
        
        logger.info(f"从({x1},{y1})拖拽到({x2},{y2})")
        return f'已从坐标 ({x1},{y1}) 拖拽到 ({x2},{y2})。'
    except Exception as e:
        logger.error(f"拖拽操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 移动鼠标工具
@mcp.tool()
def move_tool(to_loc: List[int]) -> str:
    """
    将鼠标移动到特定坐标
    参数：
    to_loc: 目标坐标 [x, y]
    """
    try:
        if len(to_loc) != 2:
            raise ValueError("to_loc必须是包含两个整数的列表 [x, y]")
        
        x, y = to_loc[0], to_loc[1]
        pg.moveTo(x, y)
        
        logger.info(f"鼠标移动到坐标({x},{y})")
        return f'已将鼠标指针移动到 ({x},{y})。'
    except Exception as e:
        logger.error(f"移动鼠标失败: {str(e)}")
        return f"错误: {str(e)}"

# 快捷键工具
@mcp.tool()
def shortcut_tool(shortcut: List[str]) -> str:
    """
    执行键盘快捷键
    参数：
    shortcut: 按键组合列表，如 ["ctrl", "c"] 表示复制
    """
    try:
        pg.hotkey(*shortcut)
        
        logger.info(f"执行快捷键: {'+'.join(shortcut)}")
        return f'已按下快捷键 {'+'.join(shortcut)}。'
    except Exception as e:
        logger.error(f"执行快捷键失败: {str(e)}")
        return f"错误: {str(e)}"

# 按键工具
@mcp.tool()
def key_tool(key: str = '') -> str:
    """
    按下单个键盘按键
    参数：
    key: 要按下的键
    """
    try:
        pg.press(key)
        
        logger.info(f"按下键: {key}")
        return f'已按下键 {key}。'
    except Exception as e:
        logger.error(f"按键操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 等待工具
@mcp.tool()
def wait_tool(duration: int) -> str:
    """
    暂停执行指定的秒数
    参数：
    duration: 等待时间（秒）
    """
    try:
        time.sleep(duration)
        
        logger.info(f"等待了{duration}秒")
        return f'已等待 {duration} 秒。'
    except Exception as e:
        logger.error(f"等待操作失败: {str(e)}")
        return f"错误: {str(e)}"

# 网页抓取工具
@mcp.tool()
def scrape_tool(url: str) -> str:
    """
    获取网页内容并转换为markdown格式
    参数：
    url: 完整的URL（包含http/https）
    """
    try:
        import requests
        from bs4 import BeautifulSoup
        
        response = requests.get(url, timeout=10)
        soup = BeautifulSoup(response.text, 'html.parser')
        
        # 提取主要文本内容
        for script in soup(['script', 'style']):
            script.decompose()
        
        text = soup.get_text(separator='\n', strip=True)
        
        logger.info(f"已抓取网页: {url}")
        return f'已抓取网页内容:\n{text[:1000]}...'
    except ImportError:
        return "错误: 缺少必要的库 requests 和 BeautifulSoup"
    except Exception as e:
        logger.error(f"网页抓取失败: {str(e)}")
        return f"错误: {str(e)}"

# 主程序入口
if __name__ == "__main__":
    logger.info("\n\n\t桌面控制服务已启动！等待调用！\n\n")
    mcp.run(transport="stdio")